import numpy as np
from skimage.metrics import structural_similarity as ssim
import argparse
import os

def read_yuv420p_frame(file, width, height, frame_idx):
    """
    读取 YUV420P 文件的某一帧
    :param file: 文件对象（已打开）
    :param width: 视频宽度
    :param height: 视频高度
    :param frame_idx: 帧索引（从0开始）
    :return: Y, U, V 分量（NumPy 数组）
    """
    # 计算一帧的字节大小
    y_size = width * height
    uv_size = (width // 2) * (height // 2)
    frame_size = y_size + 2 * uv_size  # Y + U + V
    
    # 定位到目标帧
    file.seek(frame_idx * frame_size)
    
    # 读取 Y 分量
    y_data = np.frombuffer(file.read(y_size), dtype=np.uint8)
    y = y_data.reshape((height, width))
    
    # 读取 U 分量
    u_data = np.frombuffer(file.read(uv_size), dtype=np.uint8)
    u = u_data.reshape((height // 2, width // 2))
    
    # 读取 V 分量
    v_data = np.frombuffer(file.read(uv_size), dtype=np.uint8)
    v = v_data.reshape((height // 2, width // 2))
    
    return y, u, v

def calculate_yuv_ssim(yuv_file1, yuv_file2, width, height, frame_idx):
    """
    计算两帧 YUV420P 的 SSIM（基于 Y 分量）
    """
    # 读取两帧的 Y 分量
    y1, _, _ = read_yuv420p_frame(yuv_file1, width, height, frame_idx)
    y2, _, _ = read_yuv420p_frame(yuv_file2, width, height, frame_idx)
    
    # 计算 SSIM
    score, _ = ssim(y1, y2, full=True, data_range=255)
    return score

def write_line(data,yuv_file_name):
	with open('output.txt', 'a') as file:
		file.write("input:"+yuv_file_name+"       ")
		file.write(data)
		file.write('\n')


def calculate_yuv_video_ssim(yuv_path1, yuv_path2, width, height, max_frames=None):
    """
    计算两个 YUV420P 视频的平均 SSIM
    :param yuv_path1: 第一个 YUV 文件路径
    :param yuv_path2: 第二个 YUV 文件路径
    :param width: 视频宽度
    :param height: 视频高度
    :param max_frames: 最大计算帧数（None 表示全部）
    :return: 平均 SSIM
    """
    with open(yuv_path1, 'rb') as f1, open(yuv_path2, 'rb') as f2:
        ssim_scores = []
        frame_idx = 0
        
        while True:
            try:
                # 计算当前帧 SSIM
                score = calculate_yuv_ssim(f1, f2, width, height, frame_idx)
                ssim_scores.append(score)
                frame_idx += 1
                
                # 限制最大帧数
                if max_frames and frame_idx >= max_frames:
                    break
            except Exception as e:
                print(f"计算结束或出错: {e}")
                break
        
        if not ssim_scores:
            return 0.0
        
       	avg_ssim = np.mean(ssim_scores)
        write_line(f"Average SSIM: {avg_ssim:.4f} ",yuv_path1)
        print(f"计算完成，共 {frame_idx} 帧")
        print(f"SSIM 范围: [{np.min(ssim_scores):.4f}, {np.max(ssim_scores):.4f}]")
        print(f"平均 SSIM (Y 分量): {avg_ssim:.4f}")
        return avg_ssim

if __name__ == "__main__":
    # 命令行参数解析
    parser = argparse.ArgumentParser(description="计算两个 YUV420P 视频的 SSIM")
    parser.add_argument("yuv1", help="第一个 YUV 文件路径")
    parser.add_argument("yuv2", help="第二个 YUV 文件路径")
    parser.add_argument("--width", type=int, required=True, help="视频宽度")
    parser.add_argument("--height", type=int, required=True, help="视频高度")
    parser.add_argument("--max_frames", type=int, default=None, help="最大计算帧数")
    args = parser.parse_args()
    
    # 检查文件是否存在
    if not os.path.exists(args.yuv1):
        raise FileNotFoundError(f"文件 {args.yuv1} 不存在")
    if not os.path.exists(args.yuv2):
        raise FileNotFoundError(f"文件 {args.yuv2} 不存在")
    
    # 计算 SSIM
    calculate_yuv_video_ssim(args.yuv1, args.yuv2, args.width, args.height, args.max_frames)
